from serpapi import GoogleSearch
import requests
from PIL import Image
from io import BytesIO
import imghdr
import os
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry


def download_image(image_url, save_path):
    """根据图片URL下载图片并验证合法性"""
    try:
        session = requests.Session()
        retries = Retry(total=5, backoff_factor=0.1, status_forcelist=[500, 502, 503, 504])
        session.mount('https://', HTTPAdapter(max_retries=retries))
        response = requests.get(image_url, stream=True, timeout=10, verify=False)

        # 检查是否为图片类型
        content_type = response.headers.get('Content-Type')
        if 'image' not in content_type:
            print(f"跳过非图片链接: {image_url}")
            return False

        # 将响应数据转换为字节并用PIL检查
        img_data = BytesIO(response.content)
        img = Image.open(img_data)

        # 保存图片
        img.save(save_path)

        # 验证图片文件是否合法
        if imghdr.what(save_path) is None:
            print(f"下载图片格式不正确: {image_url}")
            os.remove(save_path)
            return False

        print(f"图片已保存：{save_path}")
        return True

    except Exception as e:
        print(f"下载链接响应出错 {image_url}. 错误信息: {e}")
        return False


def search_image(query, save_directory, max_images=10, step=1):
    params = {
        "engine": "yandex_images",
        "text": query,
        "image_type": "photo",
        "api_key": "ac6f7dbdf12c1e1cb9a038ed4a3661cdc8fb87441160eb712b5f98001532f3db"
    }

    search = GoogleSearch(params)
    results = search.get_dict()
    for i in range(max_images):
        image_url = results["images_results"][i]["original"]
        if image_url.find("jpg") != -1:
            file_extension = ".jpg"
        elif image_url.find("png") != -1:
            file_extension = ".png"
        elif image_url.find("jpeg") != -1:
            file_extension = ".jpeg"
        else:
            file_extension = ".webp"
        download_image(image_url, save_directory + f"/image_{step}_{i}" + file_extension)


if __name__ == '__main__':
    search_image("coffee","img", 10, 1)
